package com.ly.utils;

import ai.onnxruntime.*;

import java.util.Map;

public class OnnxModelUtils {

    public static void printModelInputsAndOutputs(OrtSession session) {
        try {
            // 打印输入信息
            System.out.println("Model Inputs:");
            for (Map.Entry<String, NodeInfo> entry : session.getInputInfo().entrySet()) {
                NodeInfo inputInfo = entry.getValue();
                System.out.println("Name: " + inputInfo.getName());
                System.out.println("Type: " + inputInfo.getInfo());
                System.out.println();
            }

            // 打印输出信息
            System.out.println("Model Outputs:");
            for (Map.Entry<String, NodeInfo> entry : session.getOutputInfo().entrySet()) {
                NodeInfo outputInfo = entry.getValue();
                System.out.println("Name: " + outputInfo.getName());
                System.out.println("Type: " + outputInfo.getInfo());
                System.out.println();
            }

        } catch (OrtException e) {
            System.err.println("Error occurred while retrieving ONNX model information: " + e.getMessage());
            e.printStackTrace();
        }
    }

    public static void main(String[] args) {
        if (args.length != 1) {
            System.err.println("Usage: java OnnxModelUtils <model_path>");
            System.exit(1);
        }
        String modelPath = args[0];
        
        try (OrtEnvironment env = OrtEnvironment.getEnvironment();
             OrtSession session = env.createSession(modelPath, new OrtSession.SessionOptions())) {

            printModelInputsAndOutputs(session);

        } catch (OrtException e) {
            System.err.println("Error occurred while working with ONNX model: " + e.getMessage());
            e.printStackTrace();
        }
    }
}
